#include <arm_sve.h>

#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <stdexcept>

#include "gemv.h"

namespace ggml {
namespace cpu {
namespace aarch64 {

void gemv_q4_0_8x8_q8_0(int n, float *GGML_RESTRICT s,
                        [[maybe_unused]] size_t bs,
                        const void *GGML_RESTRICT vx,
                        const void *GGML_RESTRICT vy, [[maybe_unused]] int nr,
                        int nc) {
  const int qk = QK8_0;
  const int nb = n / qk;
  const int ncols_interleaved = 8;
  const int blocklen = 8;

  assert(n % qk == 0);
  assert(nc % ncols_interleaved == 0);

  float sumf[8];
  int sumi;

  const block_q8_0 *a_ptr = (const block_q8_0 *)vy;
  for (int x = 0; x < nc / ncols_interleaved; x++) {
    const block_q4_0x8 *b_ptr = (const block_q4_0x8 *)vx + (x * nb);

    for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0;
    for (int l = 0; l < nb; l++) {
      for (int k = 0; k < (qk / (2 * blocklen)); k++) {
        for (int j = 0; j < ncols_interleaved; j++) {
          sumi = 0;

          // Please fill your SVE implementation here
          throw std::runtime_error("GEMV with SVE is not implemented");

          sumf[j] += sumi * GGML_FP16_TO_FP32(b_ptr[l].d[j]) *
                     GGML_FP16_TO_FP32(a_ptr[l].d);
        }
      }
    }
    for (int j = 0; j < ncols_interleaved; j++)
      s[x * ncols_interleaved + j] = sumf[j];
  }
}

template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS,
          ggml_type PARAM_TYPE>
void gemv(int, float *, size_t, const void *, const void *, int, int);

template <>
void gemv<block_q4_0, 8, 8, GGML_TYPE_Q8_0>(int n, float *s, size_t bs,
                                            const void *vx, const void *vy,
                                            int nr, int nc) {
  gemv_q4_0_8x8_q8_0(n, s, bs, vx, vy, nr, nc);
}

}  // namespace aarch64
}  // namespace cpu
}  // namespace ggml